import os
import re
import ipaddress
from datetime import datetime
from src.core.database import DatabaseConnector
from src.models.ip import IPAddress
from src.io.file_exporter import FileExporter
from src.io.csv_importer import CsvImporter
from src.core.dbhelper import DBHelper
from src.core.ip_scanner import IPScanner
from src.core.utils import *

class IPCommand:
    """IP地址命令处理类"""
    
    def __init__(self, db_connector=None, subparsers=None):
        """初始化IP命令处理器"""
        self.db_connector = db_connector or DatabaseConnector()
        self.file_exporter = FileExporter()
        self.csv_importer = CsvImporter(self.db_connector)
        self.db_helper = DBHelper(self.db_connector)
        # 初始化IP端口扫描器，不再传入数据库连接器
        self.ip_scanner = IPScanner()
        self.arg_parser = self.setup_parser(subparsers)
    
    def setup_parser(self, subparsers):
        """设置IP地址管理命令行参数解析器"""
        # ip命令组
        ip_parser = subparsers.add_parser('ip', help='IP地址管理相关命令', description='IP地址管理相关的所有操作')
        ip_subparsers = ip_parser.add_subparsers(dest='ip_command', help='IP地址管理子命令')

        # ip help命令组
        ip_subparsers.add_parser('help', help='显示IP地址管理命令帮助')

        # ip add命令组
        ip_add_parser = ip_subparsers.add_parser('add', help='添加IP地址')
        ip_add_parser.add_argument('ip', help='IP地址')
        ip_add_parser.add_argument('--owner', help='所属单位ID/名称')

        # ip set命令组
        ip_set_parser = ip_subparsers.add_parser('set', help='设置IP地址信息')
        ip_set_parser.add_argument('identifier', help='ID/IP地址')
        ip_set_parser.add_argument('--owner', help='所属单位ID/名称')

        # ip get命令组
        ip_get_parser = ip_subparsers.add_parser('get', help='获取IP地址信息')
        ip_get_parser.add_argument('identifier', help='ID/IP地址')
        ip_get_parser.add_argument('relation_type', nargs='?', choices=['domain', 'asset'], help='关系类型')

        # ip getall命令组
        ip_getall_parser = ip_subparsers.add_parser('getall', help='获取所有IP地址列表')
        ip_getall_parser.add_argument('-o', '--output', nargs='?', const=True, help='导出文件路径，不指定则自动生成')
        ip_getall_parser.add_argument('-e', '--extension', choices=['json', 'csv'], default='csv', help='导出文件格式，默认csv')


        # ip search命令组
        ip_search_parser = ip_subparsers.add_parser('search', help='搜索IP地址')
        ip_search_parser.add_argument('keyword', help='搜索关键词')
        ip_search_parser.add_argument('-o', '--output', nargs='?', const=True, help='导出文件路径，不指定则自动生成')
        ip_search_parser.add_argument('-e', '--extension', choices=['json', 'csv'], default='csv', help='导出文件格式，默认csv')
        ip_search_parser.add_argument('-s', '--scan', action='store_true', help='是否扫描IP地址，默认不扫描')
        ip_search_parser.add_argument('-f', '--scan-filter', choices=['all', 'unscan'], default='unscan', help='扫描过滤条件')
        ip_search_parser.add_argument('-p', '--scan-port', help='扫描端口(写法例如：80,443,8000-8080)，默认为常见端口')
        ip_search_parser.add_argument('-w', '--scan-timeout', default=2, help='扫描超时时间，默认为2秒')
        ip_search_parser.add_argument('--scan-retry', default=1, help='扫描失败后重试次数，默认为1次')
        ip_search_parser.add_argument('--scan-threads', default=100, help='每个IP的扫描线程数，默认为100个线程')
        ip_search_parser.add_argument('--scan-protocol', choices=['tcp', 'udp', 'all'], default='tcp', help='扫描协议，默认为tcp')
        ip_search_parser.add_argument('--scan-speed', default=100, help='扫描速度，针对每个IP每秒请求数，默认为100')
        # 目标在线检测，不在线的不进行扫描
        ip_search_parser.add_argument('--check-alive', action='store_true', help='PIMG检测目标IP是否在线，不在线的不进行扫描，默认不检测')
        # 联网检测，避免网络中断所有探测无效
        ip_search_parser.add_argument('--network-check-ip', default='8.8.8.8', help='设置联网检测的测试IP地址，避免进行无效探测')

        # ip scan命令组
        ip_scan_parser = ip_subparsers.add_parser('scan', help='扫描IP地址端口')
        ip_scan_parser.add_argument('ip', help='ID或IP地址，多个用逗号分隔')
        ip_scan_parser.add_argument('-f', '--scan-filter', choices=['all', 'unscan'], default='unscan', help='扫描过滤条件')
        ip_scan_parser.add_argument('-p', '--scan-port', help='扫描端口(写法例如：80,443,8000-8080)，默认为常见端口')
        ip_scan_parser.add_argument('-w', '--scan-timeout', default=2, help='扫描超时时间，默认为2秒')
        ip_scan_parser.add_argument('--scan-retry', default=1, help='扫描失败后重试次数，默认为1次')
        ip_scan_parser.add_argument('--scan-threads', default=100, help='每个IP的扫描线程数，默认为100个线程')
        ip_scan_parser.add_argument('--scan-protocol', choices=['tcp', 'udp', 'all'], default='tcp', help='扫描协议，默认为tcp')
        ip_scan_parser.add_argument('--scan-speed', default=100, help='扫描速度，针对每个IP每秒请求数，默认为100')
        # 目标在线检测，不在线的不进行扫描
        ip_scan_parser.add_argument('--check-alive', action='store_true', help='PIMG检测目标IP是否在线，不在线的不进行扫描，默认不检测')
        # 联网检测，避免网络中断所有探测无效
        ip_scan_parser.add_argument('--network-check-ip', default='8.8.8.8', help='设置联网检测的测试IP地址，避免进行无效探测')
    


        # ip delete命令组
        ip_delete_parser = ip_subparsers.add_parser('delete', help='删除IP地址')
        ip_delete_parser.add_argument('identifier', help='ID/IP地址')


        # ip import命令组
        ip_import_parser = ip_subparsers.add_parser('import', help='导入IP地址数据')
        ip_import_parser.add_argument('file_path', help='CSV文件路径')
        ip_import_parser.add_argument('--ip-column', help='IP列的索引或名称，默认为列名称为"ip"的列， 或第一列', default='ip')
        ip_import_parser.add_argument('--owner-column', help='单位名称列的索引或名称，默认为列名称为"owner"的列， 或第二列', default='owner')

        return ip_parser

    def handle_command(self, args):
        """处理IP命令"""
        command = args.ip_command
        if command == 'help':
            self.show_ip_help()
        elif command == 'add':
            self.add_ip(args.ip, args.owner)
        elif command == 'set':
            self.set_ip(args.identifier, args.owner)
        elif command == 'get':
            self.get_ip(args.identifier, args.relation_type)
        elif command == 'getall':
            export_path = args.output if hasattr(args, 'export') else None
            extension = args.extension if hasattr(args, 'extension') else None
            self.get_all_ips(export_path=export_path, extension=extension)
        elif command == 'scan':
            self.scan_ip(args.ip, args.scan_filter, args.scan_port, args.scan_timeout, args.scan_retry, args.scan_threads, 
                         args.scan_protocol, args.scan_speed, args.check_alive, args.network_check_ip)
        elif command == 'search':
            export_path = args.output
            extension = args.extension
            self.search_ips(args.keyword, export_path, extension, args.scan, args.scan_filter, 
                          args.scan_port, args.scan_timeout, args.scan_retry, args.scan_threads, 
                          args.scan_protocol, args.scan_speed, args.check_alive, args.network_check_ip)
        elif command == 'delete':
            self.delete_ip(args.identifier)
        elif command == 'import':
            self.import_ips(args.file_path, args.ip_column, args.owner_column)
        else:
            print(f"未知的IP命令: {command}")
            self.show_ip_help()
    
    def show_ip_help(self):
        """显示IP地址管理命令帮助信息"""
        self.arg_parser.print_help()
    
    def _get_ip(self, identifier):
        return self.db_helper.get_ip(identifier)

    def _get_owner(self, identifier):
        """获取单位"""
        return self.db_helper.get_owner(identifier)
           
    def add_ip(self, identifier, owner_identifier=None):
        """添加新IP地址"""
        # 处理所属单位
        owner_id = None
        if owner_identifier:
            owner = self._get_owner(owner_identifier)
            if owner:
                owner_id = owner['owner_id']
            else:
                print(f"错误：未找到单位: {owner_identifier}")
                return
        
        # 添加IP地址
        try:
            ip = self.db_helper.get_or_add_ip(identifier, owner_id)
            print(f"  ID={ip['ip_id']}")
            print(f"  IP地址: {ip['ip']}")
            if owner_id:
                owner_name = owner['owner_name']
                print(f"  所属单位: {owner_name} (ID: {owner_id})")
        except Exception as e:
            print(f"添加IP地址失败: {e}")
    
    def set_ip(self, identifier, owner_identifier=None):
        """设置IP地址所属单位"""
        ip = self._get_ip(identifier)
        if not ip:
            print(f"错误：未找到IP地址: {identifier}")
            return

        # 处理所属单位
        owner_id = None
        if owner_identifier:
            owner = self._get_owner(owner_identifier)
            if owner:
                owner_id = owner['owner_id']
            else:
                print(f"错误：未找到单位: {owner_identifier}")
                return

            # 更新IP地址所属单位
            current_time = datetime.now()
            sql = "UPDATE ips SET owner_id = %s, update_time = %s WHERE ip_id = %s"
            try:
                self.db_connector.execute_query(sql, (owner_id, current_time, ip['ip_id']))
                owner_name = owner['owner_name']
                print(f"成功设置IP地址 '{ip['ip']}' 的所属单位为: {owner_name} (ID: {owner_id})")
            except Exception as e:
                print(f"设置IP地址 '{ip['ip']}' 的所属单位失败: {e}")

    def get_ip(self, ip, relation_type=None):
        """获取指定IP地址信息"""
        ip = self._get_ip(ip)
        if not ip:
            print(f"未找到IP地址: {ip}")
            return

        print("IP地址信息:")
        print(f"  ID: {ip['ip_id']}")
        print(f"  IP地址: {ip['ip']}")
        print(f"  所属单位: {ip['owner_name']} (ID: {ip['owner_id']})")
        print(f"  创建时间: {ip['create_time']}")
        print(f"  更新时间: {ip['update_time']}")
        
        print("")
        if relation_type == 'domain':
            self._show_ip_domains(ip)
        elif relation_type == 'asset':
            self._show_ip_assets(ip)
    
    def get_all_ips(self, export_path=None, extension=None):
        """获取所有IP地址列表"""
        # 查询所有IP地址
        sql = "SELECT i.*, o.owner_name FROM ips i LEFT JOIN owners o ON i.owner_id = o.owner_id ORDER BY i.ip"
        ips_result = self.db_connector.fetchall(sql)

        if len(ips_result) == 0:
            print("暂无IP地址信息！")
            return
        
        # 准备导出数据
        export_data = []
        for ip in ips_result:
            export_item = {
                'ip_id': ip['ip_id'],
                'ip': ip['ip'],
                'owner_id': ip['owner_id'],
                'create_time': ip['create_time'],
                'update_time': ip['update_time'],
                'owner_name': ip['owner_name']
            }
            # 转换datetime对象为字符串
            if 'create_time' in export_item and isinstance(export_item['create_time'], datetime):
                export_item['create_time'] = export_item['create_time'].strftime('%Y-%m-%d %H:%M:%S')
            if 'update_time' in export_item and isinstance(export_item['update_time'], datetime):
                export_item['update_time'] = export_item['update_time'].strftime('%Y-%m-%d %H:%M:%S')
            export_data.append(export_item)
        
        # 导出数据
        if export_path is not None:
            actual_path = None if export_path is True else export_path
            self.file_exporter.export_data(export_data, actual_path, 'all_ips', extension)
            return
        
        # 显示IP列表
        json_head = [
            {'name': 'ID', 'key': 'ip_id', 'width': 8},
            {'name': 'IP地址', 'key': 'ip', 'width': 15},
            {'name': '所属单位', 'key': 'owner_name', 'width': 30, 'align': 'center'},
            {'name': '扫描状态', 'key': 'scan_status', 'width': 8, 'dict': {0: '未扫描', 1: '已扫描', 2: '再扫描'}},
            {'name': '端口扫描量', 'key': 'port_scan_count', 'width': 10, 'align': 'center'},
            {'name': '最后扫描时间', 'key': 'last_scan_time', 'width': 19, 'align': 'center'},
        ]
        show_table('IP地址列表', json_head, ips_result)
    
    def scan_ip(self, ip_identifier, scan_filter='unscan', scan_port=None, scan_timeout=None, scan_retry=None, scan_threads=None, scan_protocol=None, scan_speed=None, check_alive=False, network_check_ip='8.8.8.8'):
        """扫描指定IP地址"""
        ips = self.db_helper.get_or_add_ips(ip_identifier)
        self._scan_ips(ips, scan_filter, scan_port, scan_timeout, scan_threads, scan_protocol, scan_speed, scan_retry, check_alive, network_check_ip)

    def search_ips(self, keyword, export_path=None, extension=None, scan=False, scan_filter='unscan', scan_port=None, scan_timeout=None, scan_retry=None, scan_threads=None, scan_protocol=None, scan_speed=None, check_alive=False, network_check_ip='8.8.8.8'):
        """搜索IP地址"""
        # 模糊匹配IP地址或所属单位名称
        sql = "SELECT i.*, o.owner_name FROM ips i LEFT JOIN owners o ON i.owner_id = o.owner_id WHERE i.ip LIKE %s OR o.owner_name LIKE %s ORDER BY i.ip"
        params = [f"%{keyword}%", f"%{keyword}%"]
        
        ips_result = self.db_connector.fetchall(sql, params)
        if len(ips_result) == 0:
            print(f"未找到包含 '{keyword}' 的IP地址！")
            return

        if scan:
            self._scan_ips(ips_result, scan_filter, scan_port, scan_timeout, scan_threads, scan_protocol, scan_speed, scan_retry, check_alive, network_check_ip)
            return
        
        # 准备导出数据
        export_data = []
        for ip in ips_result:
            export_item = {
                'ip_id': ip['ip_id'],
                'ip': ip['ip'],
                'owner_id': ip['owner_id'],
                'create_time': ip['create_time'],
                'update_time': ip['update_time'],
                'owner_name': ip['owner_name']
            }
            # 转换datetime对象为字符串
            if 'create_time' in export_item and isinstance(export_item['create_time'], datetime):
                export_item['create_time'] = export_item['create_time'].strftime('%Y-%m-%d %H:%M:%S')
            if 'update_time' in export_item and isinstance(export_item['update_time'], datetime):
                export_item['update_time'] = export_item['update_time'].strftime('%Y-%m-%d %H:%M:%S')
            export_data.append(export_item)
        
        # 导出数据
        if export_path is not None:
            actual_path = None if export_path is True else export_path
            self.file_exporter.export_data(export_data, actual_path, 'ip_search', extension)
            return
        
        # 显示搜索结果
        json_head = [
            {'name': 'ID', 'key': 'ip_id', 'width': 8 },
            {'name': 'IP地址', 'key': 'ip', 'width': 15 },
            {'name': '所属单位', 'key': 'owner_name', 'width': 30, 'align': 'center' },
            {'name': 'PING可达', 'key': 'is_pingable', 'width': 8, 'dict': {0: '不可达', 1: '可达'}},
            {'name': '扫描状态', 'key': 'scan_status', 'width': 8, 'dict': {0: '未扫描', 1: '已扫描', 2: '再扫描'}},
            {'name': '端口扫描量', 'key': 'port_scan_count', 'width': 10, 'align': 'center'},
            {'name': '最后扫描时间', 'key': 'last_scan_time', 'width': 19, 'align': 'center'},
        ]
        show_table('匹配的IP地址', json_head, ips_result)
        

    def delete_ip(self, ip):
        """删除IP地址"""
        ip = self._get_ip(ip)
        if not ip:
            print(f"错误：IP地址 {ip} 不存在")
            return
        
        ip_id = ip['ip_id']
        # 删除IP地址
        sql = "DELETE FROM ips WHERE ip_id = %s"
        try:
            self.db_connector.execute_query(sql, (ip_id,))
            print(f"成功删除IP地址: {ip['ip']} (ID={ip_id})")
        except Exception as e:
            print(f"删除IP地址失败: {e}")
    
    def import_ips(self, file_path, ip_column, owner_column=None):
        """从CSV文件导入IP地址数据"""
        # 验证文件是否存在
        if not os.path.exists(file_path):
            print(f"错误: 文件 {file_path} 不存在！")
            return
        
        # 检查文件扩展名
        file_ext = os.path.splitext(file_path)[1].lower()
        if file_ext != '.csv':
            print(f"错误: 只支持CSV文件格式，当前文件格式为 '{file_ext}'")
            return
        
        try:
            success_count, failed_count = self.csv_importer.import_ips_from_csv(file_path, ip_column, owner_column)
            print(f"IP地址导入完成: 成功 {success_count} 个, 失败 {failed_count} 个")
        except Exception as e:
            print(f"导入IP地址时发生错误: {str(e)}")
    
    def _show_ip_assets(self, ip):
        """显示IP地址关联的资产信息"""
        ip_address = ip['ip']
        # 查询关联的资产数量
        sql = "SELECT * FROM assets WHERE ip = %s"
        assets_result = self.db_connector.fetchall(sql, (ip_address,))
        
        if len(assets_result) == 0:
            print(f"IP地址 {ip_address} 下没有资产")
            return
        
        asset_count = len(assets_result)
        print(f"IP地址 {ip_address} 下有 {asset_count} 个资产:")
        print("{:<10} {:<20} {:<10} {:<15} {:<5} {:<3}".format("资产ID", "资产名称", "域名", "IP", "端口", "协议"))
        print("=" * 60)
        for asset in assets_result:
            ip_address = asset['ip'] if asset['ip'] else '无'
            domain_name = asset['domain'] if asset['domain'] else '无'
            print("{:<10} {:<20} {:<10} {:<15} {:<5} {:<3}".format(
                asset['asset_id'], asset['asset_name'], domain_name, ip_address, asset['port'], asset['asset_protocol']
            ))
    
    def _show_ip_domains(self, ip):
        """显示IP地址关联的域名信息"""
        ip_address = ip['ip']
        # 查询关联的域名数量
        sql = "SELECT * FROM ipdomains ido LEFT JOIN domains on ido.domain = domains.domain WHERE ip = %s"
        domains_result = self.db_connector.fetchall(sql, (ip_address,))

        if len(domains_result) == 0:
            print(f"IP地址 {ip_address} 下没有域名")
            return

        domain_count = len(domains_result)
        print(f"IP地址 {ip_address} 下有 {domain_count} 个域名:")
        print("{:<10} {:<20}".format("域名ID", "域名"))
        print("=" * 40)
        for domain in domains_result:
            domain_id = domain['domain_id'] if domain['domain_id'] else '无'
            domain_name = domain['domain'] if domain['domain'] else '无'
            print("{:<10} {:<20}".format(domain_id, domain_name))

    def _scan_ips(self, ips, scan_filter, scan_port=None, scan_timeout=None, scan_threads=None, scan_protocol=None, scan_speed=None, scan_retry=None, check_alive=False, network_check_ip='8.8.8.8'):
        """扫描IP地址端口"""
        # 按照过滤类型过滤IP地址
        filterd_ips = []
        if scan_filter == 'all':
            filterd_ips = ips
        elif scan_filter == 'unscan':  # 查询未扫描的IP地址
            filterd_ips = [ip for ip in ips if not ip.get('scan_status')]
        
        if len(filterd_ips) == 0:
            print("没有需要扫描的IP地址")
            return
        
        # 根据last_scan_time对IP进行排序，扫描时间越早的，越优先扫描
        # 处理None值的情况，确保没有扫描过的IP（last_scan_time为None）排在最前面
        filterd_ips.sort(key=lambda x: x.get('last_scan_time') if x.get('last_scan_time') is not None else datetime.min)

        # 准备扫描参数
        scan_params = {}
        
        if scan_timeout is not None:
            try:
                scan_params['timeout'] = int(scan_timeout)
            except ValueError:
                print(f"警告: 扫描超时时间 '{scan_timeout}' 不是有效的整数，将使用默认值")
        
        if scan_threads is not None:
            try:
                # 设置每个进程的线程数
                scan_params['threads_per_process'] = int(scan_threads)
            except ValueError:
                print(f"警告: 扫描线程数 '{scan_threads}' 不是有效的整数，将使用默认值")
        
        if scan_speed is not None:
            try:
                scan_params['scan_speed'] = int(scan_speed)
            except ValueError:
                print(f"警告: 扫描速度 '{scan_speed}' 不是有效的整数，将使用默认值")
        
        if scan_retry is not None:
            try:
                scan_params['retry_count'] = int(scan_retry)
            except ValueError:
                print(f"警告: 扫描重试次数 '{scan_retry}' 不是有效的整数，将使用默认值")

        # 检查是否需要检查IP是否在线
        scan_params['check_alive'] = check_alive

        if network_check_ip:
            scan_params['network_check_ip'] = network_check_ip
        
        # 一次性设置所有扫描参数并打印日志
        if scan_params:
            self.ip_scanner.set_scan_params(**scan_params)
        
        # 解析端口参数
        tcp_ports = None
        udp_ports = None
        ports = None
        if scan_port:
            ports = []
            try:
                # 解析端口范围和单个端口，端口数值只能在1-65535之间
                for part in scan_port.split(','):
                    part = part.strip()  # 去除空格
                    if not part:  # 跳过空字符串
                        continue
                    if '-' in part:
                        start, end = part.split('-')
                        ports.extend(range(max(0, int(start.strip())), min(65535, int(end.strip()) + 1)))
                    else:
                        ports.append(min(65535,max(0,int(part))))
                # 检测端口是否重复，去重
                ports = list(set(ports))
            except ValueError:
                print(f"警告: 扫描端口 '{scan_port}' 格式无效，将使用默认端口")
        
                
        # 根据协议类型设置端口
        if scan_protocol == 'all':
            tcp_ports = ports
            udp_ports = ports
        elif scan_protocol == 'tcp':
            tcp_ports = ports
            udp_ports = []
        elif scan_protocol == 'udp':
            tcp_ports = []
            udp_ports = ports
        self.ip_scanner.set_ports(tcp_ports=tcp_ports, udp_ports=udp_ports, quiet=True)

        # 扫描IP地址，这里就是引入IPScanner功能模块进行IP端口扫描的地方
        self.ip_scanner.scan_ips(filterd_ips)